program cubeflu

use basic_var
use mpi_basic_var
use timing
use mpi
#ifdef _ACCEL
use accel_lib
#endif
#ifdef _OPENMP
use omp_lib
#endif

real(myk) :: start_time,end_time,elapsed_time
real(myk) :: step_start_time,step_end_time,step_elapsed_time=0._myk
real(myk) :: io_step_start_time,io_step_end_time,io_step_elapsed_time=0._myk
real(myk) :: post_step_start_time,post_step_end_time,post_step_elapsed_time=0._myk
integer :: ik

call start_mpi()

GPU = .false.

#ifdef _ACCEL
print*,'NEW n_rank,NCOREXNODE: ',n_rank,NCOREXNODE
if(mod(n_rank,NCOREXNODE) .lt. NGPUXNODE) then
  call acc_set_device(acc_device_nvidia)
  device_kind = acc_get_device()
  print*,'Selected device kind:  ',device_kind
  num_devices = acc_get_num_devices(acc_device_nvidia)
  print*,'Number of devices available: ',num_devices
  call acc_set_device_num(mod(n_rank,NCOREXNODE),acc_device_nvidia)
  print*,'n_rank: ',n_rank,' tries to set GPU: ',mod(n_rank,NCOREXNODE)
  my_device = acc_get_device_num(acc_device_nvidia)
  print*,'n_rank: ',n_rank,' is using device: ',my_device
  print*,'Set GPU to true for rank: ',n_rank
  GPU = .true.
endif
#endif
#ifdef _OPENMP
if(.not. GPU) then
   if(openmp_activate == 1) then
     print*,'Setting OpenMP threads :',NCOREXNODE - NGPUXNODE
     call omp_set_num_threads(NCOREXNODE - NGPUXNODE) ; NCOREXNODE = NGPUXNODE + 1
!$omp parallel
!$omp master
     print*,'Set OpenMP threads :',omp_get_num_threads()
!$omp end master
!$omp end parallel

   else ; call omp_set_num_threads(1) ; endif
endif
#else
if(openmp_activate == 1) then
  print*,'Error: please compile using OpenMP or disable openmp_actiavte in mpi_gpu.dat' ; call end_mpi() ; STOP
endif
#endif

call time_init()
call size_init()
call decompose_mpi()
call allocate_storage()
print*,'MPI and Memory allocated succesfully'

call random_init()
call input_grid()
call input_grid_tot()
call ini_eval()

call read_scalars()

if(itmin.eq.0) then
   call var_init()
!   call print_sol(it)
!   call print_sol_binary()
   call write_sol_binary_pio()
   call print_vtk()
else   ! restart
!   call read_sol(it)
!   call read_sol_binary
   call read_sol_binary_pio()
   t = itmin*dt
endif

#ifdef POST_RUNTIME
if(itmin.lt.post_itmin) then
   call init_pos()
else
   call read_average_tecplot_space_and_time(1)
endif
#endif

call var_k_init()
! call rhs_init() is inside time-advancement loop
call rhs_k_init()

call compact_init()

if(GPU) then
  call update_gpu()
  print*,'GPU activated for rank : ',n_rank
endif
call print_screen()
call mpi_barrier(mpi_comm_world,ierr)
call time(start_time)
do it=itmin+1,itmax
   call time(step_start_time)
   t = t+dt
   if(n_rank.eq.0) then
     if(mod(it,itout).eq.0) then
       print*, '__________________________________________________'
       print*, 'it = ',it,' t = ',t
     endif
   endif
   !zero rhs for the next step
   if(GPU) then ; call rhs_init_acc() ;       else ; call rhs_init_omp() ;       endif
   do ik=1,s
      tk = t + alpha(ik)*dt
      if(GPU) then ; call var_k_eval_acc(ik) ;       else ; call var_k_eval_omp(ik) ;    endif
      if(GPU) then ; call update_var_k_mpi_acc() ;   else ; call update_var_k_mpi_omp() ;   endif
      if(GPU) then ; call bc_var_k_acc() ;           else ; call bc_var_k_omp() ;           endif
      if(GPU) then ; call rhs_k_eval_acc() ;         else ; call rhs_k_eval_omp() ;         endif
      if(GPU) then ; call bc_rhs_k_acc() ;           else ; call bc_rhs_k_omp() ;           endif
      if(GPU) then ; call rhs_eval_acc(ik) ;         else ; call rhs_eval_omp(ik) ;         endif
   enddo
! HARMFUL   if(GPU) then ; call bc_rhs_acc() ;            else ; call bc_rhs_omp() ;          endif
   if(GPU) then ; call var_eval_acc() ;          else ; call var_eval_omp() ;        endif
   if(GPU) then ; call update_var_mpi_acc() ;    else ; call update_var_mpi_omp() ;  endif
! USELESS if(GPU) then ; call bc_var_acc() ;            else ; call bc_var_omp() ;          endif 
   if(GPU) then ; call post_eval_acc() ;         else ; call post_eval_omp() ;       endif

   call time(step_end_time)
   step_elapsed_time = step_elapsed_time + (step_end_time - step_start_time)
   if(mod(it,itout).eq.0) then
     call time(io_step_start_time)
     if(GPU) call update_cpu()
     call print_screen()
!     call print_sol()
!     call print_sol_binary()
     call write_sol_binary_pio()
     call print_vtk()
     call time(io_step_end_time)
     io_step_elapsed_time = io_step_end_time - io_step_start_time
     if(n_rank.eq.0) then
       open(unit=1,file='iter.dat')
         write(1,*) it
         write(1,*) it + itmax-itmin
         write(1,*) itout
         write(1,*) dt
#ifdef POST_RUNTIME
         write(1,*) post_itmin  ;  write(1,*) it + itmax-itmin  ;  write(1,*) post_itout
#endif
       close(1)
     endif
     if(n_rank.eq.0) print*,'Step Elapsed time : ',step_elapsed_time/itout ; step_elapsed_time = 0._myk
     if(n_rank.eq.0) print*,'IO/Step Elapsed time : ',io_step_elapsed_time 
   endif
#ifdef POST_RUNTIME
   post_it = it
   if((it .ge. post_itmin) .and. (mod(it,post_itout).eq.0)) then
     call time(post_step_start_time)
     if(GPU) call update_cpu()
     call print_screen()

     call enlarge_mpi_space_and_time_first()
     call post_mpi_first()
     call reduce_mpi_space_and_time_first()

     call enlarge_mpi_space_and_time_second()
     call init_pos_space_average()
     call post_mpi_second()
     call reduce_mpi_space_second()
     call reduce_mpi_space_and_time_second()
     if(n_rank .eq. 0) call print_average_tecplot_space(1)

     call time(post_step_end_time)
     post_step_elapsed_time = post_step_end_time - post_step_start_time
     if(n_rank.eq.0) print*,'POST/Step Elapsed time : ',post_step_elapsed_time 
   endif
#endif
enddo

it = itmax ! otherwise it=itmax+1 
!call print_sol_binary()

#ifdef POST_RUNTIME
if(n_rank .eq. 0) call print_average_tecplot_space_and_time(1)
#endif

call time(end_time)
!print*,'ending Runge-Kutta'
elapsed_time = end_time - start_time
print*,'Elapsed time : ',elapsed_time,' GPU? ',GPU

call end_mpi()

end program cubeflu
